from datetime import datetime
import gradio as gr
import json
import requests
from string import Template
import pyaudio, wave
import os, sys
import time

now_dir = os.getcwd()
sys.path.append(now_dir)
sys.path.append(os.path.join(now_dir, "GPT_SoVITS"))

os.environ["OPENAI_API_KEY"] = "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"
os.environ["SERPAPI_API_KEY"] = "9200796cba4b2569d70549a440e2ee16c690401256613ec3cdeffc46edb47652"

session_id = "foo"

port = 5000
# 取得模型文件夹路径
config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
if os.path.exists(config_path):
    with open(config_path, "r", encoding="utf-8") as f:
        _config = json.load(f)
        port = _config.get("tts_port", 5000)
        default_batch_size = _config.get("batch_size", 10)
        is_share = _config.get("is_share", "false").lower() == "true"

default_request_url = f"http://127.0.0.1:{port}"
character_info_endpoint = f"{default_request_url}/character_list"
tts_endpoint = f"{default_request_url}/tts"
chat_endpoint = f"{default_request_url}/chat"
save_chat_endpoint = f"{default_request_url}/save_chat"
get_redis_chat_endpoint = f"{default_request_url}/redis_chat"
asr_endpoint = f"{default_request_url}/asr"
default_read_data = """{
    "method": "POST",
    "body": {
        "cha_name": "${chaName}",
        "character_emotion": "${characterEmotion}",
        "text": "${speakText}",
        "text_language": "${textLanguage}",
        "batch_size": ${batch_size},
        "speed": ${speed_factor},
        "top_k": ${topK},
        "top_p": ${topP},
        "temperature": ${temperature},
        "stream": "${stream}",
        "save_temp": "True"
    }
}"""
default_chat_data = """{
    "method": "POST",
    "body": {
        "session_id":"${sessionID}",
        "input_type":"${inputType}",
        "history":${history}
    }
}"""
default_asr_data = """{
    "method": "POST",
    "body": {
        "session_id":"${sessionID}",
        "input_audio":"${inputAudio}",
        "history":${history}
    }
}"""

global p, streamAudio
p = pyaudio.PyAudio()
streamAudio = None


def load_character_emotions(character_name, characters_and_emotions):
    emotion_options = ["default"]
    emotion_options = characters_and_emotions.get(character_name, ["default"])

    return gr.Dropdown(emotion_options, value="default")


def read_aloud(
        endpoint_data,
        text,
        cha_name,
        text_language,
        batch_size,
        speed_factor,
        top_k,
        top_p,
        temperature,
        character_emotion,
        read_mode,
        stream="False",

):
    urlencoded_text = requests.utils.quote(text)

    # 使用Template填充变量
    params = {
        "chaName": cha_name,
        "speakText": urlencoded_text,
        "textLanguage": text_language,
        "batch_size": batch_size,
        "speed_factor": speed_factor,
        "topK": top_k,
        "topP": top_p,
        "temperature": temperature,
        "characterEmotion": character_emotion,
        "stream": stream,
    }

    endpoint_data_template = Template(endpoint_data)
    filled_json_str = endpoint_data_template.substitute(**params)
    # 解析填充后的JSON字符串
    request_data = json.loads(filled_json_str)
    body = request_data["body"]
    if stream.lower() == "false":
        print(f"发送请求到{tts_endpoint}")
        # 发送POST请求
        response = requests.post(tts_endpoint, json=body)
        # 检查请求是否成功
        if response.status_code == 200:
            # 生成保存路径
            save_path = (
                f"tmp_audio/{cha_name}{datetime.now().strftime('%Y%m%d%H%M%S%f')}.wav"
            )

            # 检查保存路径是否存在
            if not os.path.exists("tmp_audio"):
                os.makedirs("tmp_audio")

            # 保存音频文件到本地
            with open(save_path, "wb") as f:
                f.write(response.content)

            # 返回给gradio
            return gr.Audio(save_path, type="filepath")
        else:
            gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
            return gr.Audio(None, type="filepath")
    else:
        if read_mode:
            global p, streamAudio
            # 打开音频流
            streamAudio = p.open(format=p.get_format_from_width(2),
                                 channels=1,
                                 rate=32000,
                                 output=True)

            response = requests.post(tts_endpoint, json=body, stream=True)
            if response.status_code == 200:
                save_path = (
                    f"tmp_audio/{cha_name}{datetime.now().strftime('%Y%m%d%H%M%S%f')}.wav"
                )

                # 音频参数
                channels = 1  # 单声道
                sampwidth = 2  # 采样位宽，2字节（16位）
                framerate = 32000  # 采样率，32000 Hz

                # 检查保存路径是否存在
                if not os.path.exists("tmp_audio"):
                    os.makedirs("tmp_audio")

                # 打开一个新的 wave 文件，准备写入
                with wave.open(save_path, 'wb') as wf:
                    wf.setnchannels(channels)  # 设置声道数
                    wf.setsampwidth(sampwidth)  # 设置采样位宽
                    wf.setframerate(framerate)  # 设置采样率
                    for data in response.iter_content(chunk_size=1024):
                        wf.writeframes(data)
                        if (streamAudio is not None) and (not streamAudio.is_stopped()):
                            streamAudio.write(data)

                # 停止和关闭流
                if streamAudio is not None:
                    streamAudio.stop_stream()
                return gr.Audio(save_path, type="filepath")
            else:
                gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
                return gr.Audio(None, type="filepath")
        else:
            return gr.Audio(None, type="filepath")


def get_bot_message(
        endpoint_data,
        session_id,
        input_type,
        history
):
    urlencoded_history = []
    history[-1][1] = "tmp"
    for pair in history:
        urlencoded_history += [[requests.utils.quote(pair[0]), requests.utils.quote(pair[1])]]

        # 使用Template填充变量
    params = {
        "sessionID": session_id,
        "inputType": input_type,
        "history": urlencoded_history
    }

    endpoint_data_template = Template(endpoint_data)
    filled_json_str = endpoint_data_template.substitute(**params)
    filled_json_str = filled_json_str.replace("'", "\"")
    # 解析填充后的JSON字符串
    request_data = json.loads(filled_json_str)
    body = request_data["body"]
    history[-1][1] = ""
    print(f"发送请求到{chat_endpoint}")
    # 发送POST请求
    response = requests.post(chat_endpoint, json=body, stream=True)
    # 检查请求是否成功
    if response.status_code == 200:

        for i, line in enumerate(response.iter_content(chunk_size=2048)):
            if i == 0:
                print(f"stream out:{time.time()}")
            if line:
                history[-1][1] += line.decode()
            yield history
    else:
        gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
        history[-1][1] += "不好意思，程序出了一些问题"
        return history


def save_bot_message(
        endpoint_data,
        session_id,
        history
):
    print(f"stream end:{time.time()}")
    urlencoded_history = []
    for pair in history:
        urlencoded_history += [[requests.utils.quote(pair[0]), requests.utils.quote(pair[1])]]

        # 使用Template填充变量
    params = {
        "sessionID": session_id,
        "inputType": "text",
        "history": urlencoded_history
    }
    endpoint_data_template = Template(endpoint_data)
    filled_json_str = endpoint_data_template.substitute(**params)
    filled_json_str = filled_json_str.replace("'", "\"")
    # 解析填充后的JSON字符串
    request_data = json.loads(filled_json_str)
    body = request_data["body"]
    print(f"发送请求到{save_chat_endpoint}")
    # 发送POST请求
    response = requests.post(save_chat_endpoint, json=body)
    # 检查请求是否成功
    if response.status_code == 200:
        return history[-1][1]
    else:
        gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
        return history[-1][1]


def get_redis_message(
        session_id
):
    # 使用Template填充变量
    params = {
        "session_id": session_id
    }
    print(f"发送请求到{get_redis_chat_endpoint}")
    # 发送GET请求
    response = requests.get(get_redis_chat_endpoint, params=params)
    # 检查请求是否成功
    if response.status_code == 200:
        redis_history = json.loads(response.content)

        return redis_history
    else:
        gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
        return []


def speech_recognition(
        endpoint_data,
        session_id,
        audio,
        history
):
    urlencoded_history = []
    history += [["录音错误，请你忽略这句话", "tmp"]]
    for pair in history:
        urlencoded_history += [[requests.utils.quote(pair[0]), requests.utils.quote(pair[1])]]
    # 使用Template填充变量
    params = {
        "sessionID": session_id,
        "inputAudio": audio,
        "history": urlencoded_history
    }
    endpoint_data_template = Template(endpoint_data)
    filled_json_str = endpoint_data_template.substitute(**params)
    filled_json_str = filled_json_str.replace("'", "\"")
    filled_json_str = filled_json_str.replace("\\", "\\\\")
    # 解析填充后的JSON字符串
    request_data = json.loads(filled_json_str)
    request_data["input_audio"] = audio
    body = request_data["body"]
    print(f"发送请求到{asr_endpoint}")
    # 发送POST请求
    response = requests.post(asr_endpoint, json=body)
    # 检查请求是否成功
    if response.status_code == 200:
        history = json.loads(response.content)
        return history
    else:
        gr.Warning(f"请求失败，状态码：{response.status_code}, 返回内容：{response.content}")
        return history


def stopAudioPlay():
    global streamAudio
    if streamAudio is not None:
        streamAudio.stop_stream()
        streamAudio = None


def get_characters_and_emotions(character_list_url):
    try:
        response = requests.get(character_list_url)
        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"请求失败，状态码：{response.status_code}")
    except:
        raise Exception("请求失败，请检查URL是否正确")


def change_character_list(
        character_list_url, cha_name="", character_emotion="default"
):
    characters_and_emotions = {}

    try:
        characters_and_emotions = get_characters_and_emotions(character_list_url)
        character_names = [i for i in characters_and_emotions]
        if len(character_names) != 0:
            if cha_name in character_names:
                character_name_value = cha_name
            else:
                character_name_value = character_names[0]
        else:
            character_name_value = ""
        emotions = characters_and_emotions.get(character_name_value, ["default"])
        emotion_value = character_emotion
        if emotion_value not in emotions:
            emotion_value = "default"
    except:
        character_names = []
        character_name_value = ""
        emotions = ["default"]
        emotion_value = "default"
        characters_and_emotions = {}
    return (
        gr.Dropdown(character_names, value=character_name_value, label="选择角色"),
        gr.Dropdown(emotions, value=emotion_value, label="情感列表", interactive=True),
        characters_and_emotions,
    )


def change_batch_size(batch_size):
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            _config = json.load(f)
        with open(config_path, "w", encoding="utf-8") as f:
            _config["batch_size"] = batch_size
            json.dump(_config, f, ensure_ascii=False, indent=4)
    except:
        pass
    return
